#!/usr/bin/env python3

import rospy
import sys
import moveit_commander
import math
from std_msgs.msg import Int8
from xarm_moveit_demo.srv import *
from geometry_msgs.msg import PoseStamped, Pose
from moveit_commander import MoveGroupCommander, PlanningSceneInterface
from moveit_msgs.msg import PlanningScene, ObjectColor, Grasp, PlaceLocation, GripperTranslation, MoveItErrorCodes
from trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint
from tf.transformations import quaternion_from_euler
from copy import deepcopy


GROUP_NAME_ARM = "xarm"
GROUP_NAME_GRIPPER = "gripper"
GRIPPER_FRAME = "gripper_centor_link"
GRIPPER_OPEN = [0.65, 0.65]
GRIPPER_GRASP = [0.1, 0.1]
GRIPPER_CLOSE = [0.0, 0.0]
GRIPPER_JOINT_NAMES = ["gripper_1_joint", "gripper_2_joint"]
GRIPPER_EFFORT = [1.0, 1.0]
REFERENCE_FRAME = "base_link"

class MoveItPickPlaceDemo:
	def __init__(self):
		moveit_commander.roscpp_initialize(sys.argv)
		rospy.init_node("moveit_pick_place_demo")
		self.scene = PlanningSceneInterface()
		rospy.sleep(1)
		self.gripper_pose_pub = rospy.Publisher("gripper_pose", PoseStamped, queue_size=10)
		self.picker_place_srv = rospy.Service("/pick_place_demo", CallPickPlaceDemo, self.call_pick_place)
		self.xarm = MoveGroupCommander(GROUP_NAME_ARM)
		self.gripper = MoveGroupCommander(GROUP_NAME_GRIPPER)
		rospy.loginfo("Pick and Place demo is ready. You can call \
				the service /pick_place_demo to test...")
		rospy.spin()

		moveit_commander.roscpp_shutdown()
		moveit_commander.os._exit(0)

	def call_pick_place(self, req):
		end_effector_link = self.xarm.get_end_effector_link()
		self.xarm.allow_replanning(True)
		self.xarm.set_pose_reference_frame(REFERENCE_FRAME)
		self.xarm.set_planning_time(5)
		max_pick_attempts = 5
		max_place_attempts = 5
		table_id = "table"
		target_id = "target"
		self.scene.remove_attached_object(GRIPPER_FRAME, target_id)
		self.scene.remove_world_object(table_id)
		self.scene.remove_world_object(target_id)
		rospy.sleep(1)
		
		self.xarm.set_named_target("Home")
		self.xarm.go()
		rospy.sleep(1)

		table_size = [1.0, 1.2, 0.01]
		table_pose = PoseStamped()
		table_pose.header.frame_id = REFERENCE_FRAME
		table_pose.pose.position.x = 0.0
		table_pose.pose.position.y = 0.0
		table_pose.pose.position.z = - table_size[2] / 2.0
		table_pose.pose.orientation.w = 1.0
		self.scene.add_box(table_id, table_pose, table_size)
		
		target_size = [0.05, 0.05, 0.23]
		target_pose = PoseStamped()
		target_pose.header.frame_id = REFERENCE_FRAME
		target_pose.pose.position.x = 0.47
		target_pose.pose.position.y = 0.0
		target_pose.pose.position.z = target_size[2] / 2.0
		target_pose.pose.orientation.w = 1.0
		self.scene.add_box(target_id, target_pose, target_size)
		rospy.sleep(1)
		
		self.xarm.set_support_surface_name(table_id)
		
		grasp_pose = target_pose
		grasp_pose.pose.position.z += 0.02
		q = quaternion_from_euler(0, 0, 0)
		grasp_pose.pose.orientation.x = q[0]
		grasp_pose.pose.orientation.y = q[1]
		grasp_pose.pose.orientation.z = q[2]
		grasp_pose.pose.orientation.w = q[3]

		pre_grasp_approach_direction = [1.0, 0, 0]
		grasps = self.make_grasps(grasp_pose, [table_id], pre_grasp_approach_direction)

		for grasp in grasps:
			self.gripper_pose_pub.publish(grasp.grasp_pose)
			rospy.sleep(0.2)
		result = None
		n_attempts = 0

		while result != MoveItErrorCodes.SUCCESS and n_attempts < max_pick_attempts:
			n_attempts += 1
			rospy.loginfo("Pick attempt: " + str(n_attempts))
			result = self.xarm.pick(target_id, grasps)
			rospy.sleep(0.2)

		place_pose = PoseStamped()
		place_pose.header.frame_id = REFERENCE_FRAME
		place_pose.pose.position.x = 0.32
		place_pose.pose.position.y = -0.32
		place_pose.pose.position.z = target_size[2] / 2.0
		q = quaternion_from_euler(0, 0, -math.pi/4)
		place_pose.pose.orientation.x = q[0]
		place_pose.pose.orientation.y = q[1]
		place_pose.pose.orientation.z = q[2]
		place_pose.pose.orientation.w = q[3]

		if result == MoveItErrorCodes.SUCCESS:
			result = None
			n_attempts = 0
			places = self.make_places(place_pose)
			while result != MoveItErrorCodes.SUCCESS and n_attempts < max_place_attempts:
				n_attempts += 1
				rospy.loginfo("Place attempt: " + str(n_attempts))
				for place in places:
					result = self.xarm.place(target_id, place)
					if result == MoveItErrorCodes.SUCCESS:
						break
				rospy.sleep(0.2)
		
			if result != MoveItErrorCodes.SUCCESS:
				rospy.loginfo("Place operation failed after " + str(n_attempts) + " attempts.")
				return CallPickPlaceDemoResponse(False)
		else:
			rospy.loginfo("Pick operation failed after " + str(n_attempts) + " attempts.")
			return CallPickPlaceDemoResponse(False)
		
		self.xarm.set_named_target("Home")
		self.xarm.go()

		self.gripper.set_joint_value_target(GRIPPER_CLOSE)
		self.gripper.go()
		rospy.sleep(1)
		self.scene.remove_attached_object(GRIPPER_FRAME, target_id)
		self.scene.remove_world_object(table_id)
		self.scene.remove_world_object(target_id)
		rospy.sleep(1)
		return CallPickPlaceDemoResponse(True)

	def make_gripper_posture(self, joint_positions):
		t = JointTrajectory()
		t.joint_names = GRIPPER_JOINT_NAMES
		tp = JointTrajectoryPoint()
		tp.positions = joint_positions
		tp.effort = GRIPPER_EFFORT
		tp.time_from_start = rospy.Duration(5)
		t.points.append(tp)
		return t

	def	make_grasps(self, initial_pose_stamped, allowed_touch_objects, pre_grasp_approach_direction):
		grasps = []
		g = Grasp()
		g.pre_grasp_posture = self.make_gripper_posture(GRIPPER_OPEN)
		g.grasp_posture = self.make_gripper_posture(GRIPPER_GRASP)
		g.pre_grasp_approach = self.make_gripper_translation(0.1, 0.12, pre_grasp_approach_direction)
		g.post_grasp_retreat = self.make_gripper_translation(0.08, 0.1, [0.0, 0.0, 1.0])
		g.grasp_pose = initial_pose_stamped
		g.allowed_touch_objects = allowed_touch_objects
		g.max_contact_force = 0
		grasps.append(deepcopy(g))
		return grasps
	
	def make_gripper_translation(self, min_dist, desired, vector):
		g = GripperTranslation()
		g.direction.vector.x = vector[0]
		g.direction.vector.y = vector[1]
		g.direction.vector.z = vector[2]
		g.direction.header.frame_id = REFERENCE_FRAME
		g.min_distance = min_dist
		g.desired_distance= desired
		return g

	def make_places(self, init_pose):
		place = PlaceLocation()
		place.place_pose = init_pose
		place.pre_place_approach.direction.header.frame_id = REFERENCE_FRAME
		place.pre_place_approach.direction.vector.z = -1.0
		place.pre_place_approach.min_distance = 0.08
		place.pre_place_approach.desired_distance = 0.1

		place.post_place_retreat.direction.header.frame_id = REFERENCE_FRAME
		place.post_place_retreat.direction.vector.z = 1.0
		place.post_place_retreat.min_distance = 0.12
		place.post_place_retreat.desired_distance = 0.15

		x_vals = [0, 0.005, 0.01, -0.005, -0.01]
		y_vals = [0, 0.005, 0.01, -0.005, -0.01]
		places = []

		for y in y_vals:
			for x in x_vals:
				place.place_pose.pose.position.x = init_pose.pose.position.x + x	
				place.place_pose.pose.position.y = init_pose.pose.position.y + y
				places.append(deepcopy(place))
		return places

	
if __name__ == "__main__":
		MoveItPickPlaceDemo()